package org.mm.proxycache.util;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

import java.util.Set;

import org.apache.mina.common.AttributeKey;
import org.apache.mina.common.IoSession;
import org.apache.mina.common.IoSessionAttributeMap;
import org.mm.proxycache.util.collections.ConcurrentNavigableMap;
import org.mm.proxycache.util.collections.ConcurrentSkipListMap;

public class SessionAttributeMap implements IoSessionAttributeMap
{
    private ConcurrentNavigableMap<String, Object> map = new ConcurrentSkipListMap<String, Object>();

    public boolean containsAttribute( IoSession session, Object key )
    {
        checkNotNull( key );
        return map.containsKey( key );
    }

    public void dispose( IoSession session ) throws Exception
    {
    }

    public Object getAttribute( IoSession session, Object key, Object defaultValue )
    {
        checkNotNull( key );

        return map.get( getKey( key ) );
    }

    @SuppressWarnings("unchecked")
    public Set<Object> getAttributeKeys( IoSession session )
    {
        return (Set) map.keySet();
    }

    public Object removeAttribute( IoSession session, Object key )
    {
        checkNotNull( key );

        return map.remove( getKey( key ) );
    }

    public boolean removeAttribute( IoSession session, Object key, Object value )
    {
        checkNotNull( key );
        checkNotNull( value );

        return map.remove( getKey( key ), value );
    }

    public boolean replaceAttribute( IoSession session, Object key, Object oldValue, Object newValue )
    {
        checkNotNull( key );
        checkNotNull( oldValue );
        checkNotNull( newValue );

        return map.replace( getKey( key ), oldValue, newValue );
    }

    public Object setAttribute( IoSession session, Object key, Object value )
    {
        checkNotNull( key );
        checkNotNull( value );

        return map.put( getKey( key ), value );
    }

    public Object setAttributeIfAbsent( IoSession session, Object key, Object value )
    {
        checkNotNull( key );
        checkNotNull( value );

        return map.putIfAbsent( getKey( key ), value );
    }

    private static String getKey( Object key )
    {
        if ( key instanceof String )
            return (String) key;
        else if ( key instanceof AttributeKey )
            return ((AttributeKey) key).toString();

        checkState( false, "Invalid key type: %s", key.getClass().getName() );
        return null;
    }

}
